use crate::DataType;

#[derive(Debug)]
pub struct DataSetConfig {
    pub src: String,
    pub dst: String,
    pub data_type: DataType, // 数据类型
    pub ratio: f64, // 训练集比例
    pub class_id: i32,
}

pub fn print_usage() {
    println!("\nUsage: img2yolo -t <train|val|test> -s <src> -d <dst> -r <ratio> -c <class_id>");
    println!("\nOptions:");
    println!("  -t, --type <train|val|test>      Data type (required)");
    println!("  -s, --src <src>                  Source directory (required)");
    println!("  -d, --dst <dst>                  Destination directory (required)");
    println!("  -r, --ratio <ratio>              Training set ratio (default: 1.0)");
    println!("  -c, --class_id <class_id>        Class ID (default: 0)");
    println!("\nExample:");
    println!("  img2yolo -t train -s ./assets -d ./dataset -r 0.2 -c 0");
}

pub fn parse(args: Vec<String>) -> DataSetConfig {
    let mut config = DataSetConfig {
        data_type: DataType::Unknown,
        src: "".parse().unwrap(),
        dst: "".parse().unwrap(),
        ratio: 1_f64,
        class_id: 0
    };
    let mut i = 1;
    while i < args.len() {
        match args[i].as_str() {
            "-t" | "--type" => {
                if i + 1 < args.len() {
                    match args[i + 1].as_str() {
                        "train" => config.data_type = DataType::Train,
                        "val" => config.data_type = DataType::Val,
                        "test" => config.data_type = DataType::Test,
                        _ => config.data_type = DataType::Unknown
                    }
                    i += 1;
                }
            }
            "-s" | "--src" => {
                if i + 1 < args.len() {
                    config.src = args[i + 1].clone();
                    i += 1;
                }
            }
            "-d" | "--dst" => {
                if i + 1 < args.len() {
                    config.dst = args[i + 1].clone();
                    i += 1;
                }
            },
            "-r" | "--ratio" => {
                if i + 1 < args.len() {
                    config.ratio = args[i + 1].parse().unwrap();
                    i += 1;
                }
            }
            "-c" | "--class_id" => {
                if i + 1 < args.len() {
                    config.class_id = args[i + 1].parse().unwrap();
                    i += 1;
                }
            }
            _ => {
                print_usage();
                std::process::exit(1);
            }
        }
        i += 1;
    }
    if config.data_type.as_str() == "unknown" || config.dst.is_empty() || config.src.is_empty() {
        print_usage();
        std::process::exit(1);
    }

    config
}